import networkx as nx
import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib的字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体为微软雅黑
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

# 定义网络结构
N = 1000  # 网络中的节点数量
p = 0.01  # Erdős-Rényi网络的连接概率
G = nx.erdos_renyi_graph(N, p)

# 初始化状态变量
S = np.ones(N)  # 易感者
V = np.zeros(N)  # 疫苗接种者
E = np.zeros(N)  # 暴露者
I = np.zeros(N)  # 感染者
R = np.zeros(N)  # 康复者
D = np.zeros(N)  # 死亡者

# 初始条件：有少量感染者
I[np.random.choice(N, 10, replace=False)] = 1

# 参数设置
beta = 0.2  # 感染率
gamma = 0.1  # 疫苗接种率
sigma = 0.0  # 疫苗接种者感染率（为0）
delta = 0.1  # 暴露者转化为感染者的速率
alpha = 0.1  # 感染者康复率
rho = 0.01  # 感染者死亡率
dt = 0.1  # 时间步长
T = 160  # 总时间


# 辅助函数：计算节点的邻居中感染者的数量
def infected_neighbors(node, G, I):
    return np.sum(I[list(G.neighbors(node))])


# SVEIRD模型状态更新函数
def update_states(S, V, E, I, R, D, G, beta, gamma, sigma, delta, alpha, rho, dt):
    # 使用列表推导来计算每个易感者被其感染邻居感染的概率
    dS = -beta * S * (np.array(
        [infected_neighbors(node, G, I) / G.degree(node) if G.degree(node) > 0 else 0 for node in range(N)])) * dt
    dV = gamma * S * dt
    dE = (beta * S * (np.array([infected_neighbors(node, G, I) / G.degree(node) if G.degree(node) > 0 else 0 for node in
                                range(N)])) - delta * E) * dt
    dI = (delta * E - (alpha + rho) * I) * dt
    dR = alpha * I * dt
    dD = rho * I * dt

    # 更新总状态
    S += dS
    V += dV
    E += dE
    I += dI
    R += dR
    D += dD

    # 强制保持状态变量在合理范围内
    S[S < 0] = 0
    V[V < 0] = 0
    E[E < 0] = 0
    I[I < 0] = 0
    R[R < 0] = 0
    D[D < 0] = 0

    # 归一化疫苗接种者数量

    return S, V, E, I, R, D





# 迭代更新状态变量
total_timesteps = int(T / dt)
time_steps = np.arange(0, T, dt)
states = np.zeros((total_timesteps, 6, N))
states[0] = np.array([S, V, E, I, R, D])

for t in range(1, total_timesteps):
    S, V, E, I, R, D = update_states(S, V, E, I, R, D, G, beta, gamma, sigma, delta, alpha, rho, dt)
    states[t] = np.array([S, V, E, I, R, D])

# 可视化结果
fig, axs = plt.subplots(2, 3, figsize=(15, 10))
for i, (ax, label) in enumerate([(axs[0, 0], '易感者'), (axs[0, 1], '疫苗接种者'),
                                 (axs[0, 2], '暴露者'), (axs[1, 0], '感染者'),
                                 (axs[1, 1], '康复者'), (axs[1, 2], '死亡者')]):
    ax.plot(time_steps, np.sum(states[:, i, :], axis=1), label=label)
    ax.set_xlabel('时间/天')
    ax.legend()

plt.tight_layout()
plt.show()

# 调控措施设置
scenarios = [
    {'label': '基准情况', 'beta': 0.2, 'gamma': 0.1},
    {'label': '提高疫苗接种率', 'beta': 0.2, 'gamma': 0.2},
    {'label': '降低感染率', 'beta': 0.1, 'gamma': 0.1},
]

# 初始化结果数组
results = {}
for index, scenario in enumerate(scenarios):  # 使用 enumerate 来迭代列表
    beta = scenario['beta']
    gamma = scenario['gamma']

    # 初始化状态变量
    S = np.ones(N)
    V = np.zeros(N)
    E = np.zeros(N)
    I = np.zeros(N)
    R = np.zeros(N)
    D = np.zeros(N)

    # 初始条件：有少量感染者
    I[np.random.choice(N, 10, replace=False)] = 1

    # 迭代更新状态变量
    states_scenario = np.zeros((total_timesteps, 6, N))
    states_scenario[0] = np.array([S, V, E, I, R, D])

    for t in range(1, total_timesteps):
        S, V, E, I, R, D = update_states(S, V, E, I, R, D, G, beta, gamma, sigma, delta, alpha, rho, dt)
        states_scenario[t] = np.array([S, V, E, I, R, D])

        # 保存结果（感染者数量）
    results[scenario['label']] = np.sum(states_scenario[:, 3, :], axis=1)


# 可视化结果
fig, ax = plt.subplots(figsize=(10, 6))

# 遍历 results 字典的键
for scenario_id, values in results.items():
    ax.plot(time_steps, values, label=scenario_id)  # 使用场景的标识符作为标签

ax.set_xlabel('时间/天')
ax.set_ylabel('感染者数量')
ax.legend()
ax.set_title('不同调控措施下传播规模和速度的对比')
plt.tight_layout()
plt.show()